package smtpserver

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/nosch/smtpoll/logger"
)

// Session 表示一个 SMTP 会话状态
type Session struct {
	conn        Connection
	state       SessionState
	heloDone    bool
	mailFrom    string
	rcptTo      []string
	dataMode    bool
	messageData []byte
	buffer      *bytes.Buffer
	handler     MessageHandler
	options     *ServerOptions
	authSource  AuthSource
	ctx         context.Context
}

// NewSession 创建新的会话
func NewSession(conn Connection, handler MessageHandler, authSource AuthSource) *Session {
	return &Session{
		conn:       conn,
		state:      NewInitialState(),
		buffer:     bytes.NewBuffer(nil),
		handler:    handler,
		authSource: authSource,
		ctx:        logger.WithTraceID(context.Background()),
	}
}

// 发送响应给客户端
func (s *Session) respond(code, message string) error {
	response := fmt.Sprintf("%s %s\r\n", code, message)
	logger.LogProtocol(s.ctx, "S: "+response,
		"state", s.state.GetName())
	_, err := s.conn.Write([]byte(response))
	if err != nil {
		logger.LogError(s.ctx, "Error sending response",
			"code", code,
			"error", err)
	}
	return err
}

// 处理客户端命令
func (s *Session) handleCommand(command string) error {
	// 去除行尾的\r\n
	command = strings.TrimRight(command, "\r\n")
	logger.LogProtocol(s.ctx, "C: "+command, "state", s.state.GetName())

	// 数据模式下的特殊处理
	if s.dataMode {
		if command == "." {
			// 结束数据输入
			s.dataMode = false
			logger.LogInfo(s.ctx, "End of data", "message_size", len(s.messageData))
			s.saveMessage()
			return s.respond(CODE_OK, "Message accepted for delivery")
		}
		// 处理点转义：如果行以单个点开头，添加一个额外的点
		if strings.HasPrefix(command, ".") {
			command = "." + command
		}
		s.messageData = append(s.messageData, []byte(command+"\r\n")...)
		return nil
	}

	return s.state.HandleCommand(s, command)
}

// 处理 HELO/EHLO 命令
func (s *Session) handleHelo(cmd, params string) error {
	if params == "" {
		return s.respond(CODE_PARAMETER_ERROR, "Syntax: "+cmd+" <domain>")
	}

	s.heloDone = true
	s.mailFrom = ""
	s.rcptTo = nil
	s.state = NewAuthenticatedState()

	// 简单响应，实际服务器可能返回支持的扩展
	if cmd == CMD_EHLO {
		capabilities := []string{
			"STARTTLS",
			"8BITMIME",
			"SMTPUTF8",
			"PIPELINING",
		}

		// 构建多行响应
		var response strings.Builder
		// 第一行：250-<domain>
		response.WriteString(fmt.Sprintf("250-%s\r\n", params))
		// 中间行：250-<capability>
		for _, cap := range capabilities[:len(capabilities)-1] {
			response.WriteString(fmt.Sprintf("250-%s\r\n", cap))
		}
		// 最后一行：250 <capability>
		response.WriteString(fmt.Sprintf("250 %s", capabilities[len(capabilities)-1]))

		return s.respond(CODE_OK, response.String())
	}

	return s.respond(CODE_OK, params)
}

// 处理 MAIL FROM 命令
func (s *Session) handleMail(params string) error {
	if !s.heloDone {
		return s.respond(CODE_BAD_SEQUENCE, "HELO/EHLO first")
	}

	if !strings.HasPrefix(strings.ToUpper(params), "FROM:") {
		return s.respond(CODE_PARAMETER_ERROR, "Syntax: MAIL FROM:<address>")
	}

	// 提取邮箱地址
	fromAddr := strings.TrimPrefix(params, "FROM:")
	fromAddr = strings.Trim(fromAddr, "<> ")

	if fromAddr == "" {
		return s.respond(CODE_PARAMETER_ERROR, "Invalid address")
	}

	// 增加事务编号
	s.ctx = logger.IncrementTransactionNo(s.ctx)

	// 重置会话状态
	s.mailFrom = fromAddr
	s.rcptTo = nil

	return s.respond(CODE_OK, "Sender ok")
}

// 处理 RCPT TO 命令
func (s *Session) handleRcpt(params string) error {
	if !s.heloDone {
		return s.respond(CODE_BAD_SEQUENCE, "HELO/EHLO first")
	}

	if s.mailFrom == "" {
		return s.respond(CODE_BAD_SEQUENCE, "MAIL FROM first")
	}

	if !strings.HasPrefix(strings.ToUpper(params), "TO:") {
		return s.respond(CODE_PARAMETER_ERROR, "Syntax: RCPT TO:<address>")
	}

	// 提取邮箱地址
	toAddr := strings.TrimPrefix(params, "TO:")
	toAddr = strings.Trim(toAddr, "<> ")

	if toAddr == "" {
		return s.respond(CODE_PARAMETER_ERROR, "Invalid address")
	}

	// 检查邮箱是否存在（这里简化处理，总是返回成功）
	s.rcptTo = append(s.rcptTo, toAddr)
	return s.respond(CODE_OK, "Recipient ok")
}

// 处理 DATA 命令
func (s *Session) handleData() error {
	if !s.heloDone {
		return s.respond(CODE_BAD_SEQUENCE, "HELO/EHLO first")
	}

	if s.mailFrom == "" {
		return s.respond(CODE_BAD_SEQUENCE, "MAIL FROM first")
	}

	if len(s.rcptTo) == 0 {
		return s.respond(CODE_BAD_SEQUENCE, "RCPT TO first")
	}

	s.dataMode = true
	s.messageData = nil
	return s.respond(CODE_START_MAIL_INPUT, "End data with <CR><LF>.<CR><LF>")
}

// 处理 RSET 命令
func (s *Session) handleRset() error {
	s.mailFrom = ""
	s.rcptTo = nil
	s.dataMode = false
	s.messageData = nil
	return s.respond(CODE_OK, "Reset state")
}

// 处理 VRFY 命令
func (s *Session) handleVrfy(_ string) error {
	// VRFY 命令通常被禁用，以防止邮箱枚举攻击
	return s.respond(CODE_COMMAND_NOT_IMP, "VRFY command not supported")
}

// 处理 NOOP 命令
func (s *Session) handleNoop() error {
	return s.respond(CODE_OK, "OK")
}

// 定义特殊错误类型
var (
	ErrClientQuit = errors.New("client quit")
)

// 处理 QUIT 命令
func (s *Session) handleQuit() error {
	s.respond(CODE_BYE, "Goodbye")
	return ErrClientQuit
}

// 处理 STARTTLS 命令
func (s *Session) handleStartTLS() error {
	// 简单实现，实际需要支持 TLS
	return s.respond(CODE_COMMAND_NOT_IMP, "STARTTLS not supported")
}

// 处理认证命令
func (s *Session) handleAuth(params string) error {
	if !s.heloDone {
		return s.respond(CODE_BAD_SEQUENCE, "HELO/EHLO first")
	}

	parts := strings.Fields(params)
	if len(parts) == 0 {
		return s.respond(CODE_PARAMETER_ERROR, "Syntax: AUTH <mechanism> [initial-response]")
	}

	mechanism := strings.ToUpper(parts[0])

	switch mechanism {
	case AUTH_PLAIN:
		return (*Session).handleAuthPlain(s, parts[1:])
	case AUTH_LOGIN:
		return s.handleAuthLogin(parts[1:])
	default:
		return s.respond(CODE_COMMAND_NOT_IMP, "Unsupported authentication mechanism")
	}
}

// handleAuthPlain 处理 PLAIN 认证
func (s *Session) handleAuthPlain(params []string) error {
	var authData string
	if len(params) > 0 {
		authData = params[0]
	} else {
		// 请求认证数据
		return s.respond(CODE_AUTH_CONTINUE, "")
	}

	// 解码认证数据
	data, err := base64.StdEncoding.DecodeString(authData)
	if err != nil {
		return s.respond(CODE_AUTH_FAILED, "Invalid authentication data")
	}

	// PLAIN 格式: \0username\0password
	parts := bytes.Split(data, []byte{0})
	if len(parts) != 3 {
		return s.respond(CODE_AUTH_FAILED, "Invalid authentication format")
	}

	username := string(parts[1])
	password := string(parts[2])

	// 使用认证源验证
	if s.authSource.Authenticate(username, password) {
		s.state = NewAuthenticatedState()
		return s.respond(CODE_AUTH_OK, "Authentication successful")
	}

	return s.respond(CODE_AUTH_FAILED, "Authentication failed")
}

// handleAuthLogin 处理 LOGIN 认证
func (s *Session) handleAuthLogin(params []string) error {
	var username, password []byte
	var err error

	if len(params) == 0 {
		// 请求用户名
		return s.respond(CODE_AUTH_CONTINUE, "VXNlcm5hbWU6") // "Username:" in base64
	}

	// 解码用户名
	username, err = base64.StdEncoding.DecodeString(params[0])
	if err != nil {
		return s.respond(CODE_AUTH_FAILED, "Invalid username format")
	}

	if len(params) == 1 {
		// 请求密码
		return s.respond(CODE_AUTH_CONTINUE, "UGFzc3dvcmQ6") // "Password:" in base64
	}

	// 解码密码
	password, err = base64.StdEncoding.DecodeString(params[1])
	if err != nil {
		return s.respond(CODE_AUTH_FAILED, "Invalid password format")
	}

	// 使用认证源验证
	if s.authSource.Authenticate(string(username), string(password)) {
		s.state = NewAuthenticatedState()
		return s.respond(CODE_AUTH_OK, "Authentication successful")
	}

	return s.respond(CODE_AUTH_FAILED, "Authentication failed")
}

// 保存邮件消息
func (s *Session) saveMessage() {
	// 创建邮件消息
	msg := &Message{
		From: s.mailFrom,
		To:   s.rcptTo,
		Data: s.messageData,
		Date: time.Now(),
	}

	logger.LogInfo(s.ctx, "Message saved",
		"from", s.mailFrom,
		"to", s.rcptTo,
		"size", len(s.messageData))

	// 使用消息处理器处理邮件
	if err := s.handler.HandleMessage(msg); err != nil {
		logger.LogError(s.ctx, "Error handling message", "error", err)
	}
}

// 设置会话状态
func (s *Session) SetState(state SessionState) {
	s.state = state
}

// 获取当前状态
func (s *Session) GetState() SessionState {
	return s.state
}

// 重置会话状态
func (s *Session) Reset() {
	s.heloDone = false
	s.mailFrom = ""
	s.rcptTo = nil
	s.dataMode = false
	s.messageData = nil
	s.state = NewInitialState()
}

// StartTLSState STARTTLS状态
type StartTLSState struct{}

// HandleCommand 处理 STARTTLS 命令
func (s *StartTLSState) HandleCommand(session *Session, command string) error {
	if command == "STARTTLS" {
		// 发送 220 响应
		if err := session.respond(CODE_READY, "Ready to start TLS"); err != nil {
			return err
		}

		// 创建 TLS 配置
		tlsConfig, err := createTLSConfig(session.options.TLSConfig)
		if err != nil {
			return fmt.Errorf("创建TLS配置失败: %v", err)
		}

		// 升级连接为 TLS
		tlsConn := tls.Server(session.conn, tlsConfig)
		if err := tlsConn.Handshake(); err != nil {
			return fmt.Errorf("TLS握手失败: %v", err)
		}

		// 更新会话连接
		session.conn = tlsConn
		session.state = &ReadyState{}

		return nil
	}

	return fmt.Errorf("在STARTTLS状态下收到未知命令: %s", command)
}

// GetName 返回状态名称
func (s *StartTLSState) GetName() string {
	return "STARTTLS"
}

// ReadyState 就绪状态
type ReadyState struct{}

// HandleCommand 处理就绪状态下的命令
func (s *ReadyState) HandleCommand(session *Session, command string) error {
	parts := strings.Fields(command)
	if len(parts) == 0 {
		return session.respond(CODE_SYNTAX_ERROR, "无效的命令")
	}

	switch strings.ToUpper(parts[0]) {
	case "HELO", "EHLO":
		return session.respond(CODE_OK, "OK")
	case "MAIL":
		return session.respond(CODE_OK, "OK")
	case "RCPT":
		return session.respond(CODE_OK, "OK")
	case "DATA":
		return session.respond(CODE_START_MAIL_INPUT, "Start mail input; end with <CRLF>.<CRLF>")
	case "QUIT":
		return session.respond(CODE_BYE, "Bye")
	default:
		return session.respond(CODE_SYNTAX_ERROR, "Unknown command")
	}
}

// GetName 返回状态名称
func (s *ReadyState) GetName() string {
	return "READY"
}

// ProcessConnection 处理连接数据
func (s *Session) ProcessConnection(reader io.Reader) error {
	// 创建带缓冲的读取器
	bufReader := bufio.NewReader(reader)

	// 发送欢迎消息
	if err := s.respond(CODE_READY, "SMTP server ready"); err != nil {
		return fmt.Errorf("failed to send welcome message: %v", err)
	}

	// 处理会话
	for {
		// 读取一行（直到遇到 CRLF）
		command, err := bufReader.ReadString('\n')
		if err != nil {
			logger.LogError(s.ctx, "Error reading from connection", "error", err)
			return fmt.Errorf("error reading from connection: %v", err)
		}

		if err := s.handleCommand(command); err != nil {
			if err == ErrClientQuit {
				// 客户端正常退出，不记录错误
				return nil
			}
			logger.LogError(s.ctx, "Error handling command", "command", command, "error", err)
			return fmt.Errorf("error handling command: %v", err)
		}
	}
}
